// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/functions/lib/aggregates/tests/utils/AggregationTestBase.h>
#include <pollux/vector/vector_fuzzer.h>

using kumo::pollux::exec::test::AssertQueryBuilder;
using namespace kumo::pollux::plan;

namespace kumo::pollux::aggregate::prestosql {
    namespace {
        class ReduceAggTest : public functions::aggregate::test::AggregationTestBase {
        protected:
            void SetUp() override {
                AggregationTestBase::SetUp();
                disableTestStreaming();
            }

            std::vector<RowVectorPtr> fuzzData() {
                VectorFuzzer::Options opts;
                opts.vectorSize = 10'000;
                opts.nullRatio = 0.1;
                VectorFuzzer fuzzer(opts, pool());

                auto rowType = ROW({{"c0", SMALLINT()}});
                std::vector<RowVectorPtr> data;
                for (auto i = 0; i < 10; ++i) {
                    data.emplace_back(fuzzer.fuzzInputRow(rowType));
                }

                return data;
            }
        };

        TEST_F(ReduceAggTest, integersGlobal) {
            auto data = make_row_vector(
                {"c0", "c1", "m"},
                {
                    make_flat_vector<int64_t>({1, 2, 3, 4, 5, 1, 1}),
                    make_nullable_flat_vector<int64_t>(
                        {1, 2, std::nullopt, 0, 1, 3, std::nullopt}),
                    make_flat_vector<bool>({true, false, true, true, false, false, true}),
                });

            // No nulls.
            int64_t product = 1 * 2 * 3 * 4 * 5 * 1 * 1;
            auto expected = make_row_vector({make_constant(product, 1)});

            testAggregations(
                {data},
                {},
                {"reduce_agg(c0, 1, (x, y) -> (x * y), (x, y) -> (x * y))"},
                {expected});

            // No nulls. With mask.
            product = 1 * 3 * 4 * 1;
            expected = make_row_vector({make_constant(product, 1)});

            testAggregations(
                {data},
                {},
                {"reduce_agg(c0, 1, (x, y) -> (x * y), (x, y) -> (x * y)) FILTER (WHERE m)"},
                {expected});

            // Some nulls.
            product = 1 * 2 * 0 * 1 * 3;
            expected = make_row_vector({make_constant(product, 1)});

            testAggregations(
                {data},
                {},
                {"reduce_agg(c1, 1, (x, y) -> (x * y), (x, y) -> (x * y))"},
                {expected});

            // Same as above, but using a sum reduction.
            int64_t sum = 1 + 2 + 3 + 4 + 5 + 1 + 1;
            expected = make_row_vector({make_constant(sum, 1)});

            testAggregations(
                {data},
                {},
                {"reduce_agg(c0, 0, (x, y) -> (x + y), (x, y) -> (x + y))"},
                {expected});

            sum = 1 + 2 + 0 + 1 + 3;
            expected = make_row_vector({make_constant(sum, 1)});

            testAggregations(
                {data},
                {},
                {"reduce_agg(c1, 0, (x, y) -> (x + y), (x, y) -> (x + y))"},
                {expected});
        }

        TEST_F(ReduceAggTest, integersGroupBy) {
            auto data = make_row_vector(
                {"k", "c0", "m"},
                {
                    make_flat_vector<int32_t>({1, 1, 2, 2, 1}),
                    make_flat_vector<int64_t>({1, 2, 3, 4, 5}),
                    make_flat_vector<bool>({true, false, false, false, true}),
                });

            auto expected = make_row_vector({
                make_flat_vector<int32_t>({1, 2}),
                make_flat_vector<int64_t>({10, 12}),
            });

            testAggregations(
                {data},
                {"k"},
                {"reduce_agg(c0, 1, (x, y) -> (x * y), (x, y) -> (x * y))"},
                {expected});

            // With mask.
            expected = make_row_vector({
                make_flat_vector<int32_t>({1, 2}),
                make_nullable_flat_vector<int64_t>({5, std::nullopt}),
            });

            testAggregations(
                {data},
                {"k"},
                {"reduce_agg(c0, 1, (x, y) -> (x * y), (x, y) -> (x * y)) FILTER (WHERE m)"},
                {expected});

            // Sum reduction.
            expected = make_row_vector({
                make_flat_vector<int32_t>({1, 2}),
                make_flat_vector<int64_t>({8, 7}),
            });

            testAggregations(
                {data},
                {"k"},
                {"reduce_agg(c0, 0, (x, y) -> (x + y), (x, y) -> (x + y))"},
                {expected});

            // With mask.
            expected = make_row_vector({
                make_flat_vector<int32_t>({1, 2}),
                make_nullable_flat_vector<int64_t>({6, std::nullopt}),
            });

            testAggregations(
                {data},
                {"k"},
                {"reduce_agg(c0, 0, (x, y) -> (x + y), (x, y) -> (x + y)) FILTER (WHERE m)"},
                {expected});
        }

        TEST_F(ReduceAggTest, arraysGlobal) {
            auto data = make_row_vector(
                {"c0", "c1", "m"},
                {
                    make_array_vector<int64_t>({
                        {1, 2, 3},
                        {1, 4, 5},
                        {},
                        {2, 6},
                        {3, 3, 7, 2},
                    }),
                    make_constant_array<int64_t>(5, {}),
                    make_flat_vector<bool>({true, false, true, false, true}),
                });

            auto expected = make_row_vector({make_array_vector<int64_t>({{7, 6, 5}})});

            testAggregations(
                {data, data, data},
                {},
                {
                    "reduce_agg(c0, c1, "
                    "(a, b) -> slice(reverse(array_sort(array_distinct(concat(a, b)))), 1, 3), "
                    "(a, b) -> slice(reverse(array_sort(array_distinct(concat(a, b)))), 1, 3))"
                },
                {expected});

            expected = make_row_vector({make_array_vector<int64_t>({{7, 3, 2}})});

            testAggregations(
                {data, data, data},
                {},
                {
                    "reduce_agg(c0, c1, "
                    "(a, b) -> slice(reverse(array_sort(array_distinct(concat(a, b)))), 1, 3), "
                    "(a, b) -> slice(reverse(array_sort(array_distinct(concat(a, b)))), 1, 3)) "
                    "FILTER (WHERE m)"
                },
                {expected});
        }

        TEST_F(ReduceAggTest, arraysGroupBy) {
            auto data = make_row_vector(
                {"k", "c0", "c1", "m"},
                {
                    make_flat_vector<int32_t>({1, 2, 1, 2, 1}),
                    make_array_vector<int64_t>({
                        {1, 2, 3},
                        {10, 20},
                        {1, 2, 4, 5},
                        {20, 30, 40},
                        {3, 1, 1, 10},
                    }),
                    make_constant_array<int64_t>(5, {}),
                    make_flat_vector<bool>({false, true, false, true, false}),
                });

            auto expected = make_row_vector({
                make_flat_vector<int32_t>({1, 2}),
                make_array_vector<int64_t>({
                    {10, 5, 4},
                    {40, 30, 20},
                }),
            });

            testAggregations(
                {data, data, data},
                {"k"},
                {
                    "reduce_agg(c0, c1, "
                    "(a, b) -> slice(reverse(array_sort(array_distinct(concat(a, b)))), 1, 3), "
                    "(a, b) -> slice(reverse(array_sort(array_distinct(concat(a, b)))), 1, 3))"
                },
                {expected});

            expected = make_row_vector({
                make_flat_vector<int32_t>({1, 2}),
                make_nullable_array_vector<int64_t>({
                    std::nullopt,
                    {{40, 30, 20}},
                }),
            });

            testAggregations(
                {data, data, data},
                {"k"},
                {
                    "reduce_agg(c0, c1, "
                    "(a, b) -> slice(reverse(array_sort(array_distinct(concat(a, b)))), 1, 3), "
                    "(a, b) -> slice(reverse(array_sort(array_distinct(concat(a, b)))), 1, 3)) "
                    "FILTER (WHERE m)"
                },
                {expected});
        }

        TEST_F(ReduceAggTest, differentInputAndCombine) {
            auto data = make_row_vector(
                {"k", "c0", "c1", "m"},
                {
                    make_flat_vector<int32_t>({1, 1, 1, 2, 2, 2, 1, 2}),
                    make_nullable_flat_vector<int32_t>(
                        {1, 2, std::nullopt, 10, std::nullopt, 20, 3, 30}),
                    make_constant_array<int32_t>(8, {}),
                    make_flat_vector<bool>(
                        {true, false, true, false, true, false, true, false}),
                });

            auto expected = make_row_vector({
                make_flat_vector<int32_t>({1, 2}),
                make_array_vector<int32_t>({{1, 2, 3}, {10, 20, 30}}),
            });

            testAggregations(
                {data, data, data},
                {"k"},
                {
                    "reduce_agg(c0, c1, "
                    "(s, x) -> array_sort(array_distinct(concat(s, array[x]))), "
                    "(s, s2) -> array_sort(array_distinct(concat(s, s2))))"
                },
                {expected});

            expected = make_row_vector({
                make_flat_vector<int32_t>({1, 2}),
                make_nullable_array_vector<int32_t>({{{1, 3}}, std::nullopt}),
            });

            testAggregations(
                {data, data, data},
                {"k"},
                {
                    "reduce_agg(c0, c1, "
                    "(s, x) -> array_sort(array_distinct(concat(s, array[x]))), "
                    "(s, s2) -> array_sort(array_distinct(concat(s, s2)))) "
                    "FILTER (WHERE m)"
                },
                {expected});
        }

        TEST_F(ReduceAggTest, fuzzGlobalSum) {
            auto data = fuzzData();
            auto plan =
                    PlanBuilder().values(data).singleAggregation({}, {"sum(c0)"}).planNode();

            auto sumResults = AssertQueryBuilder(plan).copyResults(pool());

            testAggregations(
                data,
                {},
                {"reduce_agg(c0, 0, (x, y) -> (x + y), (x, y) -> (x + y))"},
                {},
                {sumResults});
        }

        TEST_F(ReduceAggTest, fuzzGroupBySum) {
            auto data = fuzzData();
            auto plan = PlanBuilder()
                    .values(data)
                    .project({"c0 % 1234 as key", "c0"})
                    .singleAggregation({"key"}, {"sum(c0)"})
                    .planNode();

            auto sumResults = AssertQueryBuilder(plan).copyResults(pool());

            testAggregations(
                [&](PlanBuilder &builder) {
                    builder.values(data).project({"c0 % 1234 as key", "c0"});
                },
                {"key"},
                {"reduce_agg(c0, 0, (x, y) -> (x + y), (x, y) -> (x + y))"},
                {},
                [&](auto &builder) { return builder.assertResults(sumResults); });
        }

        TEST_F(ReduceAggTest, fuzzGlobalAvg) {
            auto data = fuzzData();
            auto plan =
                    PlanBuilder().values(data).singleAggregation({}, {"avg(c0)"}).planNode();

            auto avgResults = AssertQueryBuilder(plan).copyResults(pool());

            testAggregations(
                [&](PlanBuilder &builder) {
                    builder.values(data).project({
                        "c0",
                        "cast(row_constructor(0, 0) as row(sum double, count bigint)) as c1",
                    });
                },
                {},
                {
                    "reduce_agg(c0, c1, "
                    "(s, x) -> (row_constructor(s.sum + x, s.count + 1)), "
                    "(s, s2) -> (row_constructor(s.sum + s2.sum, s.count + s2.count)))"
                },
                {"a0.sum / cast(a0.count as double)"},
                [&](auto &builder) { return builder.assertResults(avgResults); });
        }

        TEST_F(ReduceAggTest, fuzzGroupByAvg) {
            auto data = fuzzData();
            auto plan = PlanBuilder()
                    .values(data)
                    .project({"c0 % 1234 as key", "c0"})
                    .singleAggregation({"key"}, {"avg(c0)"})
                    .planNode();

            auto avgResults = AssertQueryBuilder(plan).copyResults(pool());

            testAggregations(
                [&](PlanBuilder &builder) {
                    builder.values(data).project({
                        "c0",
                        "cast(row_constructor(0, 0) as row(sum double, count bigint)) as c1",
                        "c0 % 1234 as key",
                    });
                },
                {"key"},
                {
                    "reduce_agg(c0, c1, "
                    "(s, x) -> (row_constructor(s.sum + x, s.count + 1)), "
                    "(s, s2) -> (row_constructor(s.sum + s2.sum, s.count + s2.count)))"
                },
                {"key", "a0.sum / cast(a0.count as double)"},
                [&](auto &builder) { return builder.assertResults(avgResults); });
        }
    } // namespace
} // namespace kumo::pollux::aggregate::prestosql
